import numpy as np
import pandas as pd
import os
import sys
import pickle
import json
import time
from collections import defaultdict
from pytz import timezone
from datetime import datetime
from sys import getsizeof as getsize
import networkx as nx


def compare_edges(edge1, edge2):
    if len(edge1) != len(edge2):
        return -1
    for i in range(len(edge1)):
        e1 = edge1[i]
        e2 = edge2[i]
        if e1[0] != e2[0] or e1[1] != e2[1]:
            return 1
    return 0


def get_span_tag_value(span, key):
    for tag in span['tags']:
        if tag['key'] == key:
            return tag['value']
    return None


def get_endpoint(span):
    http_url = get_span_tag_value(span, 'http.method')
    http_method = get_span_tag_value(span, 'http.url')
    if http_url is None or http_method is None:
        return None
    return http_method + " " + http_url


def get_root_span(trace):
    for span in trace['spans']:
        if len(span['references']) == 0:
            return span
    return None


def get_proc_svc_map(trace):
    proc_svc_map = {}
    for p, v in trace['processes'].items():
        proc_svc_map[p] = v['serviceName']
    return proc_svc_map


def get_node_name(span, service_name):
    operation = span['operationName']
    return service_name + ":" + operation


def get_node_list(trace):
    node_list = []
    node_cntmap = {}  # 因为存在重复
    spans = sorted(trace['spans'], key=lambda x: x['startTime'])
    proc_svc_map = get_proc_svc_map(trace)
    for span in spans:
        node = get_node_name(span, proc_svc_map[span['processID']])
        cnt = 0
        if node in node_cntmap:
            cnt = node_cntmap[node]
            node_cntmap[node] = cnt + 1
        else:
            node_cntmap[node] = 1
        node_list.append(node + "_" + str(cnt))
    return node_list


def get_trace_metainfo(trace):
    
    valid = True
    trace_id = trace['traceID']
    num_span = len(trace['spans'])
    num_service = len(trace['processes'])
    
    process_map = get_proc_svc_map(trace)
    root_span = get_root_span(trace)
    if root_span is None:
        return {
            'trace_id': trace_id,
            'num_span': num_span,
            'num_service': num_service,
            # 'processes': processes,
            'valid': False,
        }
    
    root_service = None
    root_flags = None
    start_time = None
    duration = None
    root_status_code = None
    
    if root_span is not None:
        root_p = root_span['processID']
        root_service = process_map[root_p]
        root_flags = root_span['flags']
        start_time = root_span['startTime']
        duration = root_span['duration']
        root_status_code = get_span_tag_value(root_span, 'http.status_code')
        if root_status_code is None:
            valid = False
    else:
        valid = False
    
    endpoint = get_endpoint(root_span)
    if endpoint is None:
        valid = False
    
    start_time_str = datetime.fromtimestamp(start_time/1000000).strftime("%Y-%m-%d %H:%M:%S")
    duration = duration / 1000
    
    return {
        'trace_id': trace_id,
        'num_span': num_span,
        'num_service': num_service,
        # 'processes': processes,
        'root_service': root_service,
        'root_flags': root_flags,
        'start_time': start_time_str,
        'duration': str(duration) + " ms",
        'root_status_code': root_status_code,
        'endpoint': endpoint,
        'valid': valid,
    }


def get_span_map(trace):
    span_map = {}
    node_list = get_node_list(trace)
    
    spans = sorted(trace['spans'], key=lambda x: x['startTime'])
    for i, span in enumerate(spans):
        span_id = span['spanID']
        if span['references'] is None or len(span['references']) == 0:
            pid = None
        else:
            pid = span['references'][0]['spanID']
        start_time = span['startTime']
        duration = span['duration']
        node_name = node_list[i]
        span_map[span_id] = {
            'pid': pid,
            'start_time': start_time,
            'duration': duration,
            'node_name': node_name,
        }
    return span_map


def construct_graph(trace, directed=False, nginx_included=False):
    node_list = get_node_list(trace)
    node_list.sort()
    if node_list is None or len(node_list) == 0:
        return False, "invalid node_list"
    
    span_map = get_span_map(trace)
    if span_map is None or len(span_map) == 0 or len(span_map) != len(trace['spans']):
        return False, "invalid span_map"
    
    if not nginx_included:
        node_list = list(filter(lambda x: '/api' not in x, node_list))
    
    node_idx_map = {}
    idx_node_map = {}
    for idx, node in enumerate(node_list):
        node_idx_map[node] = idx
        idx_node_map[idx] = node
    
    downs_list_map = {}
    for span_id, span in span_map.items():
        pid = span['pid']
        if pid is not None and pid not in downs_list_map:
            downs_list_map[pid] = []
        if span_id not in downs_list_map:
            downs_list_map[span_id] = []
        if pid is not None:
            downs_list_map[pid].append(span_id)
    
    for span_id, span in span_map.items():
        pid = span['pid']
        if pid is None:  # root span
            continue
        if pid not in span_map:
            return False, "The upstream span {} of {} disappear".format(pid, span_id)
    
    # [total_time, rpc_time, local_time, net_time, out_degree]
    node_feats = np.zeros((len(node_list), 5))
    edge_list = []
    
    for span_id, span in span_map.items():
        node_name = span['node_name']
        
        if node_name not in node_list:
            continue
        
        idx = node_idx_map[node_name]
        pid = span['pid']
        duration = span['duration']
        
        # construct features
        duration_up = 0
        if pid is not None:  # not root span
            upper_span = span_map[pid]
            duration_up = upper_span['duration']
            
            # construct edges
            upper_node_name = upper_span['node_name']
            if upper_node_name in node_list:
                upper_idx = node_idx_map[upper_node_name]
                edge_list.append([upper_idx, idx])
                if not directed:
                    edge_list.append([idx, upper_idx])
            
        down_ids = downs_list_map[span_id]
        down_duration_sum = 0
        for child_span_id in down_ids:
            child_span = span_map[child_span_id]
            down_duration_sum += child_span['duration']
        
        total_time = duration
        rpc_time = down_duration_sum
        local_time = total_time - rpc_time
        net_time = max(0, duration_up - duration)
        out_degree = len(down_ids)

        node_feats[idx] = [total_time, rpc_time, local_time, net_time, out_degree]
    
    edge_list.sort(key=lambda x: (x[0], x[1]))
    return True, (node_feats, edge_list, idx_node_map)


def get_diff_node_pairs(edge, target_edge):
    diff_pairs = {}
    for i in range(len(edge)):
        e1 = edge[i]
        e2 = target_edge[i]
        if e1[0] != e2[0]:
            diff_pairs[min(e1[0], e2[0])] = max(e1[0], e2[0])
        if e1[1] != e2[1]:
            diff_pairs[min(e1[1], e2[1])] = max(e1[1], e2[1])
    return diff_pairs


def repermute_feature(feature, diff_pairs):
    for p1, p2 in diff_pairs.items():
        feature[[p1,p2],:] = feature[[p2,p1],:]

    
def get_trace_metric_window(trace_data, metric_all_data, time_index, 
                            endpoint_node_map, service_list, window_size=24):
    total_len = metric_all_data.shape[1]
    
    trace_feats, mt = trace_data
    endpoint = mt['endpoint']
    start_time = int(datetime.strptime(mt['start_time'], "%Y-%m-%d %H:%M:%S").timestamp())
    idx_node_map = endpoint_node_map[endpoint]
    
    # align metric window
    target_timestamp = start_time - 28800
    inc = 0
    while target_timestamp not in time_index and inc < 5:
        target_timestamp += 1
        inc += 1

    if inc >= 5:
        return None
    
    target_index = time_index[target_timestamp]
    if target_index + window_size - 1 >= total_len:
        return None
    
    metric_window = metric_all_data[:, target_index:target_index+window_size, :]
    
    # align services
    metric_service_idx_map = {s:i for i,s in enumerate(service_list)}
    trace_service_list = [node.split(':')[0] for idx, node in idx_node_map.items()]
    trace_idx_list = [metric_service_idx_map[s] for s in trace_service_list]
    
    return metric_window[trace_idx_list]